package middleware

import (
	"bp/conf"
	"bp/internal/util"
	"context"
	"crypto/rsa"
	"fmt"
	"github.com/dgrijalva/jwt-go"
	"go.uber.org/zap"
	"net/http"
	"sync"
)

const (
	bilinTokenHeadKey = "B-Token"
)

var (
	bilinPubKey  *rsa.PublicKey
	bilinPubOnce sync.Once
)

type (
	bilinClaims struct {
		jwt.StandardClaims
		*util.BilinToken
	}
)

// WithBilinToken 提取请求的token
func WithBilinToken(slog *zap.SugaredLogger, vars *conf.VarsSection) func(http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			ctx := r.Context()
			tokenStr := util.ExtractToken(r.Header.Get(bilinTokenHeadKey))
			if tokenStr == "" {
				http.Error(w, http.StatusText(http.StatusUnauthorized), http.StatusUnauthorized)
				return
			}
			tokenPtr, err := VerifyBilinToken(tokenStr, vars.BilinPubKey)
			if err != nil {
				slog.Errorf("验证token err:%#v", err.Error())
				http.Error(w, "验证token失败", http.StatusUnauthorized)
				return
			}
			slog.Infof("token got")
			ctx = context.WithValue(ctx, util.BilinTokenKey{}, tokenPtr)
			r = r.WithContext(ctx)

			next.ServeHTTP(w, r)
		})
	}
}

// LoadBilinPubKey 加载公钥
func LoadBilinPubKey(key string) func() {
	return func() {
		bilinPubKey, _ = jwt.ParseRSAPublicKeyFromPEM([]byte(key))
	}
}

// VerifyBilinToken 比邻后台 token 校验
func VerifyBilinToken(tokenStr, key string) (btoken *util.BilinToken, err error) {
	bilinPubOnce.Do(LoadBilinPubKey(key))

	parse := &jwt.Parser{
		ValidMethods:         nil,
		UseJSONNumber:        false,
		SkipClaimsValidation: true,
	}
	token, err := parse.ParseWithClaims(tokenStr, &bilinClaims{}, func(token *jwt.Token) (i interface{}, e error) {
		return bilinPubKey, nil
	})
	if err != nil {
		return
	}
	claims := token.Claims.(*bilinClaims)
	err = claims.Valid()
	if err != nil {
		return
	}
	btoken = claims.BilinToken
	if btoken == nil {
		err = fmt.Errorf("jwt payload 为空")
		return
	}
	return
}
